#!/bin/bash

# 🚀 AI Agent 一键部署脚本
# 适用于服务器: 47.83.185.154
# 用户: root
# 密码: HongKong@2024!

set -e  # 遇到错误立即退出

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 显示欢迎信息
show_welcome() {
    clear
    echo -e "${BLUE}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║                    🚀 AI Agent 一键部署脚本                    ║"
    echo "║                                                              ║"
    echo "║  服务器: 47.83.185.154                                        ║"
    echo "║  用户: root                                                   ║"
    echo "║  目标: 部署AI Agent后端服务                                    ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    echo
}

# 检查系统环境
check_environment() {
    log_info "检查系统环境..."
    
    # 检查操作系统
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        log_info "操作系统: $NAME $VERSION"
    fi
    
    # 检查Docker
    if ! command -v docker &> /dev/null; then
        log_error "Docker未安装，请先安装Docker"
        exit 1
    fi
    log_success "Docker已安装: $(docker --version)"
    
    # 检查Docker Compose
    if ! command -v docker-compose &> /dev/null; then
        log_error "Docker Compose未安装，请先安装Docker Compose"
        exit 1
    fi
    log_success "Docker Compose已安装: $(docker-compose --version)"
    
    # 检查Docker服务状态
    if ! systemctl is-active --quiet docker; then
        log_warning "Docker服务未启动，正在启动..."
        sudo systemctl start docker
        sudo systemctl enable docker
    fi
    log_success "Docker服务运行正常"
    
    # 检查端口占用
    if netstat -tlnp | grep -q :8080; then
        log_warning "端口8080已被占用"
        netstat -tlnp | grep :8080
    fi
    
    if netstat -tlnp | grep -q :3306; then
        log_warning "端口3306已被占用"
        netstat -tlnp | grep :3306
    fi
    
    # 检查磁盘空间
    DISK_USAGE=$(df / | awk 'NR==2 {print $5}' | sed 's/%//')
    if [ $DISK_USAGE -gt 80 ]; then
        log_warning "磁盘使用率较高: ${DISK_USAGE}%"
    else
        log_success "磁盘空间充足: ${DISK_USAGE}% 已使用"
    fi
}

# 创建项目目录结构
setup_directories() {
    log_info "创建项目目录结构..."
    
    PROJECT_DIR="/root/ai-agent"
    
    # 如果目录已存在，询问是否覆盖
    if [ -d "$PROJECT_DIR" ]; then
        log_warning "项目目录已存在: $PROJECT_DIR"
        read -p "是否删除现有目录并重新创建? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            rm -rf "$PROJECT_DIR"
            log_info "已删除现有目录"
        else
            log_info "使用现有目录"
        fi
    fi
    
    # 创建目录结构
    mkdir -p "$PROJECT_DIR"/{logs,data,config,backend}
    cd "$PROJECT_DIR"
    
    log_success "项目目录创建完成: $PROJECT_DIR"
    ls -la
}

# 创建配置文件
create_config_files() {
    log_info "创建配置文件..."
    
    # 创建 .env 文件
    cat > .env << 'EOF'
# AI Agent Backend Configuration
# 数据库配置
DB_TYPE=mysql
DB_HOST=mysql
DB_PORT=3306
DB_USER=aiagent
DB_PASSWORD=aiagent123
DB_NAME=aiagent

# 服务器配置
SERVER_PORT=8080
SERVER_HOST=0.0.0.0

# API安全配置
API_KEY=hk-test-secure-key-2024
JWT_SECRET=hk-jwt-secret-2024

# 日志配置
LOG_LEVEL=info
LOG_FILE=/app/logs/aiagent.log

# Agent配置
AGENT_TIMEOUT=30s
MAX_AGENTS=100

# 监控配置
METRICS_ENABLED=true
METRICS_PORT=9090
EOF
    log_success "✅ .env 文件创建完成"
    
    # 创建 docker-compose.yml 文件
    cat > docker-compose.yml << 'EOF'
version: '3.8'

services:
  ai-agent-backend:
    image: golang:1.21-alpine
    container_name: ai-agent-backend
    working_dir: /app
    ports:
      - "8080:8080"
      - "9090:9090"
    environment:
      - DB_TYPE=mysql
      - DB_HOST=mysql
      - DB_PORT=3306
      - DB_USER=aiagent
      - DB_PASSWORD=aiagent123
      - DB_NAME=aiagent
      - API_KEY=hk-test-secure-key-2024
      - SERVER_PORT=8080
      - LOG_LEVEL=info
      - CGO_ENABLED=1
    volumes:
      - ./backend:/app
      - ./logs:/app/logs
      - ./data:/app/data
    command: >
      sh -c "apk add --no-cache gcc musl-dev curl &&
             go mod tidy &&
             go build -o aiagent . &&
             ./aiagent"
    depends_on:
      mysql:
        condition: service_healthy
    networks:
      - ai-agent-network
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8080/health"]
      interval: 30s
      timeout: 10s
      retries: 3
      start_period: 60s

  mysql:
    image: mysql:8.0
    container_name: ai-agent-mysql
    environment:
      MYSQL_ROOT_PASSWORD: root123
      MYSQL_DATABASE: aiagent
      MYSQL_USER: aiagent
      MYSQL_PASSWORD: aiagent123
    ports:
      - "3306:3306"
    volumes:
      - mysql_data:/var/lib/mysql
      - ./init.sql:/docker-entrypoint-initdb.d/init.sql
    networks:
      - ai-agent-network
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "mysqladmin", "ping", "-h", "localhost"]
      interval: 10s
      timeout: 5s
      retries: 5
      start_period: 30s

volumes:
  mysql_data:
    driver: local

networks:
  ai-agent-network:
    driver: bridge
EOF
    log_success "✅ docker-compose.yml 文件创建完成"
    
    # 创建数据库初始化脚本
    cat > init.sql << 'EOF'
-- AI Agent Database Schema
CREATE DATABASE IF NOT EXISTS aiagent;
USE aiagent;

-- Agents table
CREATE TABLE IF NOT EXISTS agents (
    id VARCHAR(255) PRIMARY KEY,
    name VARCHAR(255) NOT NULL,
    status ENUM('online', 'offline', 'error') DEFAULT 'offline',
    last_heartbeat TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    ip_address VARCHAR(45),
    os_info TEXT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
);

-- Network data table
CREATE TABLE IF NOT EXISTS network_data (
    id BIGINT AUTO_INCREMENT PRIMARY KEY,
    agent_id VARCHAR(255) NOT NULL,
    interface_name VARCHAR(255),
    bytes_sent BIGINT,
    bytes_recv BIGINT,
    packets_sent BIGINT,
    packets_recv BIGINT,
    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (agent_id) REFERENCES agents(id) ON DELETE CASCADE
);

-- Process data table
CREATE TABLE IF NOT EXISTS process_data (
    id BIGINT AUTO_INCREMENT PRIMARY KEY,
    agent_id VARCHAR(255) NOT NULL,
    pid INT,
    name VARCHAR(255),
    cpu_percent DECIMAL(5,2),
    memory_percent DECIMAL(5,2),
    memory_rss BIGINT,
    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (agent_id) REFERENCES agents(id) ON DELETE CASCADE
);

-- Commands table
CREATE TABLE IF NOT EXISTS commands (
    id BIGINT AUTO_INCREMENT PRIMARY KEY,
    agent_id VARCHAR(255) NOT NULL,
    command TEXT NOT NULL,
    status ENUM('pending', 'executing', 'completed', 'failed') DEFAULT 'pending',
    result TEXT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    executed_at TIMESTAMP NULL,
    completed_at TIMESTAMP NULL,
    FOREIGN KEY (agent_id) REFERENCES agents(id) ON DELETE CASCADE
);

-- Create indexes for better performance
CREATE INDEX idx_agents_status ON agents(status);
CREATE INDEX idx_agents_last_heartbeat ON agents(last_heartbeat);
CREATE INDEX idx_network_agent_timestamp ON network_data(agent_id, timestamp);
CREATE INDEX idx_process_agent_timestamp ON process_data(agent_id, timestamp);
CREATE INDEX idx_commands_agent_status ON commands(agent_id, status);

-- Insert test data
INSERT INTO agents (id, name, status, ip_address, os_info) VALUES 
('test-agent-001', 'Test Agent 1', 'offline', '192.168.1.100', 'Windows 11'),
('test-agent-002', 'Test Agent 2', 'offline', '192.168.1.101', 'Ubuntu 22.04');
EOF
    log_success "✅ init.sql 文件创建完成"
}

# 创建后端代码
create_backend_code() {
    log_info "创建后端代码..."
    
    cd backend
    
    # 创建 go.mod 文件
    cat > go.mod << 'EOF'
module ai-agent-backend

go 1.21

require (
    github.com/gin-gonic/gin v1.9.1
    github.com/go-sql-driver/mysql v1.7.1
    github.com/golang-jwt/jwt/v5 v5.0.0
    gorm.io/driver/mysql v1.5.2
    gorm.io/gorm v1.25.5
)
EOF
    
    # 创建 main.go 文件
    cat > main.go << 'EOF'
package main

import (
    "fmt"
    "log"
    "net/http"
    "os"
    "time"

    "github.com/gin-gonic/gin"
    "gorm.io/driver/mysql"
    "gorm.io/gorm"
)

type Agent struct {
    ID            string    `json:"id" gorm:"primaryKey"`
    Name          string    `json:"name"`
    Status        string    `json:"status"`
    LastHeartbeat time.Time `json:"last_heartbeat"`
    IPAddress     string    `json:"ip_address"`
    OSInfo        string    `json:"os_info"`
    CreatedAt     time.Time `json:"created_at"`
    UpdatedAt     time.Time `json:"updated_at"`
}

type NetworkData struct {
    ID            uint      `json:"id" gorm:"primaryKey"`
    AgentID       string    `json:"agent_id"`
    InterfaceName string    `json:"interface_name"`
    BytesSent     int64     `json:"bytes_sent"`
    BytesRecv     int64     `json:"bytes_recv"`
    PacketsSent   int64     `json:"packets_sent"`
    PacketsRecv   int64     `json:"packets_recv"`
    Timestamp     time.Time `json:"timestamp"`
}

type ProcessData struct {
    ID            uint      `json:"id" gorm:"primaryKey"`
    AgentID       string    `json:"agent_id"`
    PID           int       `json:"pid"`
    Name          string    `json:"name"`
    CPUPercent    float64   `json:"cpu_percent"`
    MemoryPercent float64   `json:"memory_percent"`
    MemoryRSS     int64     `json:"memory_rss"`
    Timestamp     time.Time `json:"timestamp"`
}

var db *gorm.DB

func main() {
    // 连接数据库
    dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
        getEnv("DB_USER", "aiagent"),
        getEnv("DB_PASSWORD", "aiagent123"),
        getEnv("DB_HOST", "mysql"),
        getEnv("DB_PORT", "3306"),
        getEnv("DB_NAME", "aiagent"),
    )

    var err error
    for i := 0; i < 30; i++ {
        db, err = gorm.Open(mysql.Open(dsn), &gorm.Config{})
        if err == nil {
            break
        }
        log.Printf("Failed to connect to database (attempt %d/30): %v", i+1, err)
        time.Sleep(2 * time.Second)
    }
    
    if err != nil {
        log.Fatal("Failed to connect to database after 30 attempts:", err)
    }
    
    log.Println("Successfully connected to database")

    // 设置Gin模式
    gin.SetMode(gin.ReleaseMode)
    r := gin.Default()

    // 健康检查端点
    r.GET("/health", func(c *gin.Context) {
        c.JSON(http.StatusOK, gin.H{
            "status": "healthy",
            "time":   time.Now(),
            "version": "1.0.0",
        })
    })

    // API路由
    api := r.Group("/api")
    api.Use(authMiddleware())
    {
        api.GET("/agents", getAgents)
        api.POST("/agents", createAgent)
        api.POST("/agents/:id/heartbeat", heartbeat)
        api.POST("/agents/:id/network", uploadNetworkData)
        api.POST("/agents/:id/process", uploadProcessData)
        api.GET("/agents/:id/network", getNetworkData)
        api.GET("/agents/:id/process", getProcessData)
    }

    port := getEnv("SERVER_PORT", "8080")
    log.Printf("Server starting on port %s", port)
    log.Fatal(r.Run(":" + port))
}

func authMiddleware() gin.HandlerFunc {
    return func(c *gin.Context) {
        token := c.GetHeader("Authorization")
        expectedToken := "Bearer " + getEnv("API_KEY", "hk-test-secure-key-2024")
        
        if token != expectedToken {
            c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
            c.Abort()
            return
        }
        c.Next()
    }
}

func getAgents(c *gin.Context) {
    var agents []Agent
    db.Find(&agents)
    c.JSON(http.StatusOK, agents)
}

func createAgent(c *gin.Context) {
    var agent Agent
    if err := c.ShouldBindJSON(&agent); err != nil {
        c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
        return
    }
    
    agent.CreatedAt = time.Now()
    agent.UpdatedAt = time.Now()
    agent.Status = "online"
    
    db.Create(&agent)
    c.JSON(http.StatusCreated, agent)
}

func heartbeat(c *gin.Context) {
    agentID := c.Param("id")
    
    var agent Agent
    if err := db.First(&agent, "id = ?", agentID).Error; err != nil {
        c.JSON(http.StatusNotFound, gin.H{"error": "Agent not found"})
        return
    }
    
    agent.LastHeartbeat = time.Now()
    agent.Status = "online"
    db.Save(&agent)
    
    c.JSON(http.StatusOK, gin.H{"message": "Heartbeat received"})
}

func uploadNetworkData(c *gin.Context) {
    agentID := c.Param("id")
    
    var networkData []NetworkData
    if err := c.ShouldBindJSON(&networkData); err != nil {
        c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
        return
    }
    
    for i := range networkData {
        networkData[i].AgentID = agentID
        networkData[i].Timestamp = time.Now()
    }
    
    db.Create(&networkData)
    c.JSON(http.StatusOK, gin.H{"message": "Network data uploaded", "count": len(networkData)})
}

func uploadProcessData(c *gin.Context) {
    agentID := c.Param("id")
    
    var processData []ProcessData
    if err := c.ShouldBindJSON(&processData); err != nil {
        c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
        return
    }
    
    for i := range processData {
        processData[i].AgentID = agentID
        processData[i].Timestamp = time.Now()
    }
    
    db.Create(&processData)
    c.JSON(http.StatusOK, gin.H{"message": "Process data uploaded", "count": len(processData)})
}

func getNetworkData(c *gin.Context) {
    agentID := c.Param("id")
    
    var networkData []NetworkData
    db.Where("agent_id = ?", agentID).Order("timestamp desc").Limit(100).Find(&networkData)
    c.JSON(http.StatusOK, networkData)
}

func getProcessData(c *gin.Context) {
    agentID := c.Param("id")
    
    var processData []ProcessData
    db.Where("agent_id = ?", agentID).Order("timestamp desc").Limit(100).Find(&processData)
    c.JSON(http.StatusOK, processData)
}

func getEnv(key, defaultValue string) string {
    if value := os.Getenv(key); value != "" {
        return value
    }
    return defaultValue
}
EOF
    
    cd ..
    log_success "✅ 后端代码创建完成"
}

# 启动服务
start_services() {
    log_info "启动服务..."
    
    # 停止可能存在的旧容器
    docker-compose down 2>/dev/null || true
    
    # 启动MySQL服务
    log_info "启动MySQL服务..."
    docker-compose up -d mysql
    
    # 等待MySQL启动
    log_info "等待MySQL启动完成..."
    for i in {1..30}; do
        if docker-compose exec mysql mysqladmin ping -h localhost --silent; then
            log_success "MySQL启动成功"
            break
        fi
        echo -n "."
        sleep 2
    done
    echo
    
    # 启动后端服务
    log_info "启动AI Agent后端服务..."
    docker-compose up -d ai-agent-backend
    
    # 等待后端服务启动
    log_info "等待后端服务启动完成..."
    for i in {1..60}; do
        if curl -s http://localhost:8080/health > /dev/null 2>&1; then
            log_success "后端服务启动成功"
            break
        fi
        echo -n "."
        sleep 2
    done
    echo
    
    # 显示服务状态
    log_info "检查服务状态:"
    docker-compose ps
}

# 配置防火墙
setup_firewall() {
    log_info "配置防火墙..."
    
    # 检查ufw是否安装
    if command -v ufw &> /dev/null; then
        sudo ufw allow 8080/tcp
        sudo ufw allow 3306/tcp
        sudo ufw allow 9090/tcp
        sudo ufw reload
        log_success "防火墙配置完成"
    else
        log_warning "ufw未安装，跳过防火墙配置"
    fi
}

# 验证部署
verify_deployment() {
    log_info "验证部署..."
    
    # 健康检查
    echo "1. 健康检查:"
    if curl -s http://localhost:8080/health | grep -q "healthy"; then
        log_success "✅ 健康检查通过"
    else
        log_error "❌ 健康检查失败"
        return 1
    fi
    
    # API认证测试
    echo "2. API认证测试:"
    if curl -s -H "Authorization: Bearer hk-test-secure-key-2024" http://localhost:8080/api/agents > /dev/null; then
        log_success "✅ API认证通过"
    else
        log_error "❌ API认证失败"
        return 1
    fi
    
    # 数据库连接测试
    echo "3. 数据库连接测试:"
    if docker exec ai-agent-mysql mysql -u aiagent -paiagent123 -e "SHOW TABLES;" aiagent > /dev/null 2>&1; then
        log_success "✅ 数据库连接正常"
    else
        log_error "❌ 数据库连接失败"
        return 1
    fi
    
    # 外部访问测试
    echo "4. 外部访问测试:"
    if curl -s http://47.83.185.154:8080/health > /dev/null 2>&1; then
        log_success "✅ 外部访问正常"
    else
        log_warning "⚠️ 外部访问可能有问题，请检查防火墙设置"
    fi
    
    return 0
}

# 显示部署结果
show_deployment_result() {
    echo
    echo -e "${GREEN}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║                    🎉 部署完成！                              ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    
    echo -e "${BLUE}📋 服务信息:${NC}"
    echo "   🌐 服务器地址: http://47.83.185.154:8080"
    echo "   🔑 API密钥: hk-test-secure-key-2024"
    echo "   📊 监控地址: http://47.83.185.154:9090"
    echo "   🗄️ 数据库地址: 47.83.185.154:3306"
    echo
    
    echo -e "${BLUE}🔧 常用管理命令:${NC}"
    echo "   查看服务状态: docker-compose ps"
    echo "   查看日志: docker-compose logs <service-name>"
    echo "   重启服务: docker-compose restart <service-name>"
    echo "   停止所有服务: docker-compose down"
    echo "   启动所有服务: docker-compose up -d"
    echo
    
    echo -e "${BLUE}🧪 测试命令:${NC}"
    echo "   健康检查: curl http://47.83.185.154:8080/health"
    echo "   API测试: curl -H 'Authorization: Bearer hk-test-secure-key-2024' http://47.83.185.154:8080/api/agents"
    echo
    
    echo -e "${BLUE}📁 项目目录:${NC}"
    echo "   项目路径: /root/ai-agent"
    echo "   配置文件: /root/ai-agent/.env"
    echo "   日志目录: /root/ai-agent/logs"
    echo "   数据目录: /root/ai-agent/data"
    echo
    
    echo -e "${YELLOW}🎯 下一步:${NC}"
    echo "   1. 在本地配置客户端连接到此服务器"
    echo "   2. 使用提供的API密钥进行认证"
    echo "   3. 开始发送Agent数据到服务器"
    echo
}

# 主函数
main() {
    show_welcome
    
    log_info "开始AI Agent一键部署..."
    
    # 执行部署步骤
    check_environment
    setup_directories
    create_config_files
    create_backend_code
    start_services
    setup_firewall
    
    # 验证部署
    if verify_deployment; then
        show_deployment_result
        log_success "🎉 AI Agent部署成功完成！"
        exit 0
    else
        log_error "❌ 部署验证失败，请检查日志"
        echo "查看日志命令: docker-compose logs"
        exit 1
    fi
}

# 错误处理
trap 'log_error "部署过程中发生错误，请检查上面的错误信息"' ERR

# 运行主函数
main "$@"